import time

from exception_hyd.big_model_exception import BigModelException
import json
import requests
from typing import List, Dict, Any


# 创建一个Message类型
Message = List[Dict[str, str]]


class QianfanBigModel:
    """
    创建一个千帆大模型类
    """

    __private_api_tokens = ""
    __private_secret_tokens = ""

    def __init__(self, api_tokens: str = "", secret_tokens: str = ""):
        if not (api_tokens and secret_tokens):
            raise BigModelException("api_tokens or secret_tokens is empty")

        if not (isinstance(api_tokens, str) and isinstance(secret_tokens, str)):
            raise BigModelException("the type of api_tokens and secret_tokens are 'str'")

        self.__private_api_tokens = api_tokens
        self.__private_secret_tokens = secret_tokens

    def __get_access_tokens(self) -> str:
        """
        私有函数, 使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key

        :return: access_tokens
        """

        url = f"https://aip.baidubce.com/oauth/2.0/token?" \
              f"grant_type=client_credentials&" \
              f"client_id={ self.__private_api_tokens }&" \
              f"client_secret={ self.__private_secret_tokens}"

        payload = json.dumps("")
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload)
        return str(response.json().get("access_token"))

    def brnie_bot_turbo(self, message: Message, tp: float = 0.95, system_template: str = "") -> str:
        """
        brnie_bot_turbo模型调用方法

        :param message: 消息列表
        :param tp: 温度
        :param system_template: 模型人设
        :return: str
        """

        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?" \
              "access_token=" + self.__get_access_tokens()

        payload = json.dumps({
            "messages": message,
            "temperature": tp,
            "system": system_template
        })
        headers = {
            'Content-Type': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload).text
        response_data = json.loads(response)["result"]
        return response_data

    def brnie_bot_turbo_stream(self, message, tp: float = 0.95, system_template: str = "") -> Any:
        """
        brnie_bot_turbo模型流式响应调用方法

        :param message:
        :param tp:
        :param system_template:
        :return:
        """

        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?" \
              "access_token=" + self.__get_access_tokens()

        payload = json.dumps({
            "messages": message,
            "temperature": tp,
            "system": system_template,
            "stream": True  # 流式响应
        })
        headers = {
            'Content-Type': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload, stream=True)

        for line in response.iter_lines():
            line = line.decode("utf-8")
            if line:
                # 找到第一个冒号并提取后面的字符串
                colon_index = line.find(':')
                json_str = line[colon_index + 1:].strip()
                data = json.loads(json_str)["result"]
                if data:
                    for i in data:
                        time.sleep(0.03)
                        yield i
